/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.common.io;

import java.io.IOException;
import java.util.Arrays;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.KeyFormat;
import org.sss.util.BytesComparator;

/**
 * Formatter for System Key.
 */
public abstract class KeyFormatter {
  @SuppressWarnings("unused")
  private static Logger logger = LoggerFactory.getLogger(KeyFormatter.class);

  public static final int  KEY_GROUPID_LENGTH  = 16;
  public static final int  KEY_SUBHASH_LENGTH = 1;
  public static final int  KEY_PREFIX_LENGTH  = 1/* keyformat */ + KEY_GROUPID_LENGTH + KEY_SUBHASH_LENGTH;
  public static final byte SENTINEL           = (byte) 0xFF;

  public static final KeyFormatter userKeyPriorityFormatter = new UserKeyPriorityFormatter();
  public static final KeyFormatter threadPriorityFormatter  = new ThreadPriorityFormatter();
  public static final KeyFormatter processingOrderFormatter = new ProcessingOrderFormatter();

  private final int headerLength;
  private final int footerLength;
  private final KeyFormat format;
  private final Comparator<byte[]> comparator = new Comparator<byte[]>() {
    @Override
    public int compare(byte[] lhs, byte[] rhs) {
      int lsubhash = lhs[1 + KEY_GROUPID_LENGTH] & 0xFF;
      int rsubhash = rhs[1 + KEY_GROUPID_LENGTH] & 0xFF;

      if (lsubhash < rsubhash) return -1;
      if (lsubhash > rsubhash) return  1;

      return BytesComparator.compare(lhs, headerLength, lhs.length - headerLength,
                                     rhs, headerLength, rhs.length - headerLength);
    }
  };

  private static class UserKeyPriorityFormatter extends KeyFormatter {
    public UserKeyPriorityFormatter() {
      super(KeyFormat.USERKEY, KEY_PREFIX_LENGTH, 4 + 4 + 8);
    }

    @Override
    public void write(byte[] b, int start, int sid, int tid, long counter, byte[] rawKey) throws IOException {
      System.arraycopy(rawKey, 0, b, start, rawKey.length);
      start += rawKey.length;
      start = writeInt(b, start, sid);
      start = writeInt(b, start, tid);
      start = writeLong(b, start, counter);
    }
  }

  private static class ThreadPriorityFormatter extends KeyFormatter {
    public ThreadPriorityFormatter() {
      super(KeyFormat.THREAD, KEY_PREFIX_LENGTH + 4 + 4, 8);
    }

    @Override
    public void write(byte[] b, int start, int sid, int tid, long counter, byte[] rawKey) throws IOException {
      start = writeInt(b, start, sid);
      start = writeInt(b, start, tid);
      System.arraycopy(rawKey, 0, b, start, rawKey.length);
      start += rawKey.length;
      start = writeLong(b, start, counter);
    }
  }

  private static class ProcessingOrderFormatter extends KeyFormatter {
    public ProcessingOrderFormatter() {
      super(KeyFormat.ORDER, KEY_PREFIX_LENGTH + 4 + 4 + 8, 0);
    }

    @Override
    public void write(byte[] b, int start, int sid, int tid, long counter, byte[] rawKey) throws IOException {
      start = writeInt(b, start, sid);
      start = writeInt(b, start, tid);
      start = writeLong(b, start, counter);
      System.arraycopy(rawKey, 0, b, start, rawKey.length);
    }
  }

  /**
   * .
   *
   * @param format       KeyFormat
   * @param headerLength header length(sets in subclass).
   * @param footerLength footer length(sets in subclass).
   */
  public KeyFormatter(KeyFormat format, int headerLength, int footerLength) {
    this.format       = format;
    this.headerLength = headerLength;
    this.footerLength = footerLength;
  }

  /**
   * Generates System Key.
   */
  public byte[] generate(GroupID outputGID, int subHash, int sid, int tid, long counter, byte[] rawKey) throws IOException {
    // TODO: error, negative value(subHash, sid, tid, counter)
    byte[] b = new byte[getKeyLength(rawKey)];

    int i = writePrefix(b, outputGID, subHash);
    write(b, i, sid, tid, counter, rawKey);

    return b;
  }

  public KeyFormat getFormat() {
    return format;
  }

  public abstract void write(byte[] b, int start, int sid, int tid, long counter, byte[] rawKey) throws IOException;
  public Comparator<byte[]> getComparator() {
    return comparator;
  }

  private int getKeyLength(byte[] rawKey) {
    return headerLength + rawKey.length + footerLength;
  }

  public static byte[] getGroupIDPrefix(GroupID inputGID) throws IOException {

    byte[] b = new byte[1 + KEY_GROUPID_LENGTH];
    writePrefix(b, inputGID);
    return b;
  }

  public static byte[] getGroupIDSentinel(GroupID inputGID) throws IOException {
    return bytesInc(getGroupIDPrefix(inputGID));
  }

  public static byte[] getPrefix(GroupID inputGID, int subTaskID) throws IOException {
    byte[] b = new byte[KEY_PREFIX_LENGTH];
    writePrefix(b, inputGID, subTaskID);
    return b;
  }

  private static int writePrefix(byte[] b,
                                 GroupID inputGID) throws IOException {
    int pos = 0;
    int nslot = inputGID.getNSlot();
    if (nslot != 0 && (nslot & nslot - 1) != 0) {
      throw new IOException("Invalid number of slot, number is not two exponentiations.");
    }
    int pow = 0;
    while (nslot > 1) {
      nslot >>= 1;
      pow++;
    }
    if (pow >= 0xF) {
      throw new IOException("Invalid number of slot, too large: " + pow);
    }

    //  b   d   nslot     key format
    // [.] [.] [. . . .] [. .]
    b[pos++] = (byte)(
      ((inputGID.isBroadcast()?1:0)   << 7) |
      ((inputGID.isDestructive()?1:0) << 6) |
      ((pow & 0x0F)                   << 2) |
      (inputGID.getKeyFormat().getByte()));
    pos = writeLong(b, pos, inputGID.getMostSignificantBits());
    pos = writeLong(b, pos, inputGID.getLeastSignificantBits());
    return pos;
  }

  private static int writePrefix(byte[] b,
                                 GroupID inputGID,
                                 int subTaskID) throws IOException {
    int pos = writePrefix(b, inputGID);
    b[pos++] = (byte)(subTaskID & 0xFF); // Uses least 8 bit.
    return pos;
  }

  public static byte[] getSentinel(GroupID inputGID, int subTaskID) throws IOException {
    return bytesInc(getPrefix(inputGID, subTaskID));
  }

  public byte[] getUserKey(byte[] systemKey) {
    return Arrays.copyOfRange(systemKey, headerLength, systemKey.length - footerLength);
  }

  public static GroupID getGroupID(byte[] systemKey) {
    long m = readLong(systemKey, 1);
    long l = readLong(systemKey, 1 + 8);

    KeyFormat format = null;
    byte fmtpart = (byte)systemKey[0];
    byte fmt = (byte)(fmtpart & 0x03);
    for (KeyFormat f: KeyFormat.values()) {
      if (f.getByte() == fmt) {
        format = f;
        break;
      }
    }
    if (format == null) {
      throw new IllegalArgumentException("Invalid keyformat byte: " + fmtpart);
    }
    boolean broadcast   = ((fmtpart>> 7) & 0x1) != 0;
    boolean destructive = ((fmtpart>> 6) & 0x1) != 0;
    int nSlot = 1 << ((fmtpart >> 2) & 0x0F);
    return GroupID.create(format, m, l, broadcast, destructive, nSlot);
  }

  /**
   * .
   * This uses only ThreadPriorityFormatter.
   */
  public byte[] getWorkerIDSentinel(byte[] systemKey) {
    return bytesInc(Arrays.copyOfRange(systemKey, 0, headerLength));
  }
  public byte[] getWorkerID(byte[] systemKey) {
    return Arrays.copyOfRange(systemKey, 0, headerLength);
  }

  public boolean rawKeyEquals(byte[] systemKey, byte[] userKey) {
    int lhsLen = systemKey.length - headerLength - footerLength;

    if (lhsLen != userKey.length) {
      return false;
    }

    for (int i = 0;i < userKey.length;++i) {
      if (systemKey[headerLength + i] != userKey[i]) {
        return false;
      }
    }
    return true;
  }

  private static byte[] bytesInc(byte[] x) {
    int i = x.length;

    while (i > 0) {
      x[--i]++;
      if (x[i] != 0) break;
    }
    if (i == 0 && x[0] == 0) {
      // TODO: Uses appropriate exception.
      throw new RuntimeException("Overflow.");
    }

    return x;
  }

  // TODO: Move to util
  static String bytesToString(byte[] b) {
    StringBuilder sb = new StringBuilder();

    for (int i = 0;i < b.length;++i) {
      sb.append(String.format("%02x", b[i] & 0xFF));
      if (i % 4 == 3) {
      sb.append(" ");
      }
    }
    return sb.toString();
  }

  static public KeyFormatter get(KeyFormat format) {
    if (!keyFormatterMap.containsKey(format)) {
      throw new IllegalArgumentException("Unkown key format \"" + format + "\".");
    }
    return keyFormatterMap.get(format);
  }

  static Map<KeyFormat, KeyFormatter> keyFormatterMap = new HashMap<KeyFormat, KeyFormatter>();
  static {
    keyFormatterMap.put(userKeyPriorityFormatter.getFormat(), userKeyPriorityFormatter);
    keyFormatterMap.put(threadPriorityFormatter.getFormat(),  threadPriorityFormatter);
    keyFormatterMap.put(processingOrderFormatter.getFormat(), processingOrderFormatter);
  }

  private static int writeInt(byte[] b, int start, int value) {
    b[start++] = (byte)((value >>> (8 * 3)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 2)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 1)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 0)) & 0xFF);

    return start;
  }

  private static int writeLong(byte[] b, int start, long value) {
    b[start++] = (byte)((value >>> (8 * 7)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 6)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 5)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 4)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 3)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 2)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 1)) & 0xFF);
    b[start++] = (byte)((value >>> (8 * 0)) & 0xFF);

    return start;
  }

  @SuppressWarnings("unused")
  private static int readInt(byte[] b, int p) {
    return ( b[p]             << (8 * 3)) + // with sign bit.
           ((b[p + 1] & 0xFF) << (8 * 2)) +
           ((b[p + 2] & 0xFF) << (8 * 1)) +
           ((b[p + 3] & 0xFF) << (8 * 0));
  }

  private static long readLong(byte[] b, int p) {
    return ((long) b[p]             << (8 * 7)) + // with sign bit
           ((long)(b[p + 1] & 0xFF) << (8 * 6)) +
           ((long)(b[p + 2] & 0xFF) << (8 * 5)) +
           ((long)(b[p + 3] & 0xFF) << (8 * 4)) +
           ((long)(b[p + 4] & 0xFF) << (8 * 3)) +
           (      (b[p + 5] & 0xFF) << (8 * 2)) +
           (      (b[p + 6] & 0xFF) << (8 * 1)) +
           (      (b[p + 7] & 0xFF) << (8 * 0));
  }
}
